import mysql
import requests
import time
import random
import json
from typing import Dict, Union
from mysql.connector import Error

from applications import BaseConfig
from crawler.region_codes import PROVINCE_CODES



#产量
class AgriculturalYieldCrawler:
    def __init__(self):
        self.base_url = "https://data.stats.gov.cn/easyquery.htm"
        self.zb_mapping = {
            "A0D0Q01": "粮食产量(万吨)",
            "A0D0Q02": "夏收粮食产量(万吨)",
            "A0D0Q03": "秋粮产量(万吨)",
            "A0D0Q04": "谷物产量(万吨)",
            "A0D0Q05": "稻谷产量(万吨)",
            "A0D0Q06": "早稻产量(万吨)",
            "A0D0Q07": "中稻和一季晚稻产量(万吨)",
            "A0D0Q08": "双季晚稻产量(万吨)",
            "A0D0Q09": "小麦产量(万吨)",
            "A0D0Q0A": "冬小麦产量(万吨)",
            "A0D0Q0B": "春小麦产量(万吨)",
            "A0D0Q0C": "玉米产量(万吨)",
            "A0D0Q0D": "谷子产量(万吨)",
            "A0D0Q0E": "高粱产量(万吨)",
            "A0D0Q0F": "其他谷物产量(万吨)",
            "A0D0Q0G": "大麦产量(万吨)",
            "A0D0Q0H": "豆类产量(万吨)",
            "A0D0Q0I": "绿豆产量(万吨)",
            "A0D0Q0J": "红小豆产量(万吨)",
            "A0D0Q0K": "大豆产量(万吨)",
            "A0D0Q0L": "薯类产量(万吨)",
            "A0D0Q0M": "马铃薯产量(万吨)",
            "A0D0Q0N": "棉花产量(万吨)",
            "A0D0Q0O": "油料产量(万吨)",
            "A0D0Q0P": "花生产量(万吨)",
            "A0D0Q0Q": "油菜籽产量(万吨)",
            "A0D0Q0R": "芝麻产量(万吨)",
            "A0D0Q0S": "葵花籽产量(万吨)",
            "A0D0Q0T": "胡麻籽产量(万吨)",
            "A0D0Q0U": "糖料产量(万吨)",
            "A0D0Q0V": "甘蔗产量(万吨)",
            "A0D0Q0W": "甜菜产量(万吨)",
            "A0D0Q0X": "烟叶产量(万吨)",
            "A0D0Q0Y": "烤烟产量(万吨)",
            "A0D0Q14": "蔬菜产量(万吨)"
        }

    def _generate_headers(self) -> Dict:
        """生成随机请求头"""
        return {
            "User-Agent": random.choice([
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
                "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15"
            ]),
            "Accept": "application/json",
            "Referer": "https://data.stats.gov.cn/"
        }

    def _parse_years(self, years: Union[int, str]) -> str:
        """解析年份参数"""
        if isinstance(years, int):
            current_year = time.localtime().tm_year
            return f"{current_year - years}-{current_year}"
        return str(years)

    def get_data(self, region_code: str, years: Union[int, str] = 5) -> Dict:
        """
        获取农业产量数据
        :param region_code: 6位行政区代码，如'110000'
        :param years: 年份设置，支持：
            - int: 最近N年（包含今年）
            - str: 年份范围，如'2010-2020'
        """
        # 参数验证
        if not (len(region_code) == 6 and region_code.isdigit()):
            raise ValueError("行政区代码应为6位数字")

        # 构建时间参数
        time_range = self._parse_years(years)

        # 构造动态参数
        params = {
            "m": "QueryData",
            "dbcode": "fsnd",
            "rowcode": "zb",
            "colcode": "sj",
            "wds": json.dumps([{"wdcode": "reg", "valuecode": region_code}]),
            "dfwds": json.dumps([
                {"wdcode": "zb", "valuecode": "A0D0Q"},  # 固定农业指标
                {"wdcode": "sj", "valuecode": time_range}
            ]),
            "k1": str(int(time.time() * 1000)),  # 动态时间戳
            "h": "1"
        }

        try:
            response = requests.get(
                self.base_url,
                headers=self._generate_headers(),
                params=params,
                timeout=10
            )
            response.raise_for_status()
            return self._parse_response(response.json())
        except Exception as e:
            return {"error": str(e)}

    def _parse_response(self, data: Dict) -> Dict:
        """解析API响应"""
        result = {name: {} for name in self.zb_mapping.values()}

        for node in data.get("returndata", {}).get("datanodes", []):
            if not node["data"].get("hasdata", False):
                continue

            # 解析指标代码和年份
            code_parts = node["code"].split(".")
            zb_code = code_parts[1].split("_")[0]
            year = code_parts[-1]

            if indicator := self.zb_mapping.get(zb_code):
                result[indicator][year] = node["data"]["strdata"]

        return result


class DatabaseManager:
    def __init__(self, config):
        self.config = config
        self.create_table()

    def create_connection(self):
        """创建数据库连接"""
        try:
            return mysql.connector.connect(**self.config)
        except Error as e:
            print(f"数据库连接失败: {e}")
            return None

    def create_table(self):
        """创建数据表"""
        create_table_query = """
        CREATE TABLE IF NOT EXISTS agricultural_yield (
            id INT AUTO_INCREMENT PRIMARY KEY,
            province_code VARCHAR(6) NOT NULL,
            province_name VARCHAR(20) NOT NULL,
            year VARCHAR(4) NOT NULL,
            indicator VARCHAR(50) NOT NULL,
            value FLOAT,
            update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            UNIQUE KEY unique_record (province_code, year, indicator)
        )
        """
        conn = self.create_connection()
        if conn:
            try:
                cursor = conn.cursor()
                cursor.execute(create_table_query)
                conn.commit()
            except Error as e:
                print(f"创建表失败: {e}")
            finally:
                if conn.is_connected():
                    cursor.close()
                    conn.close()

    def save_data(self, province_code, province_name, data):
        """保存数据到数据库"""
        insert_query = """
        INSERT INTO agricultural_yield 
        (province_code, province_name, year, indicator, value)
        VALUES (%s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
        value = VALUES(value),
        update_time = CURRENT_TIMESTAMP
        """

        conn = self.create_connection()
        if not conn:
            return

        try:
            cursor = conn.cursor()
            records = []

            for indicator, years_data in data.items():
                for year, value in years_data.items():
                    try:
                        num_value = float(value) if value else None
                        records.append((
                            province_code,
                            province_name,
                            year,
                            indicator,
                            num_value
                        ))
                    except ValueError:
                        continue

            # 批量插入数据
            if records:
                cursor.executemany(insert_query, records)
                conn.commit()
                print(f"已保存 {len(records)} 条数据: {province_name}")

        except Error as e:
            print(f"数据保存失败: {e}")
            conn.rollback()
        finally:
            if conn.is_connected():
                cursor.close()
                conn.close()


def main():
    # 数据库配置
    db_config = {
        'host': BaseConfig.DB_HOST,
        'user': BaseConfig.DB_USER,
        'password': BaseConfig.DB_PASSWORD,
        'database': BaseConfig.DB_NAME,
        'port': BaseConfig.DB_PORT
    }

    # 初始化组件
    crawler = AgriculturalYieldCrawler()
    db_manager = DatabaseManager(db_config)

    # 遍历所有省份
    for province_name, province_code in PROVINCE_CODES.items():
        print(f"正在获取数据: {province_name}")

        try:
            # 获取数据（这里获取最近20年数据）
            data = crawler.get_data(province_code, 20)

            # 过滤空数据
            if not data or 'error' in data:
                print(f"数据获取失败: {province_name}")
                continue

            # 保存到数据库
            db_manager.save_data(province_code, province_name, data)

            # 添加随机延时防止被封（3-8秒）
            time.sleep(random.uniform(3, 8))

        except Exception as e:
            print(f"处理 {province_name} 时发生异常: {str(e)}")
            continue

# 使用示例
if __name__ == "__main__":
    main()
